from __future__ import annotations

from abc import ABCMeta, abstractmethod
from copy import deepcopy
from collections import Counter, namedtuple
import glob
import re
import itertools
from pathlib import Path
from typing import Literal, Any

from src.preprocess_v2 import *
from src.preprocess_v2 import Dict, Sequence
from util import *
from util import Sequence

class Functional:
    '''函数式编程'''
    def retry(self, consumer:Callable[..., Any], producer:Callable[..., Any], 
              flag:Any=None, freq:int=1, **shared_params) -> Any:
        '''重试'''
        result = consumer(**shared_params)
        if result == flag:
            while freq:
                producer(**shared_params)
                result = consumer(**shared_params)
                if result != flag:
                    break
                freq -= 1
        return result

class ExitContext:
    '''退出上下文'''
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type:  # 因为异常情况退出
            self.exc_clear()
        else:
            self.no_exc_clear()
        self.clear()
            
    def clear(self):
        pass
    
    def exc_clear(self):
        pass

    def no_exc_clear(self):
        pass

class Singleton:
    '''单类'''
    _instance = None
    def __new__(cls, *args, **kwargs):
        if not cls._instance:
            cls._instance = super().__new__(cls)
        return cls._instance
    
    def register(self, key:str, value:Any=None, default_value:Any=None) -> None:
        '''注册变量'''
        if default_value is not None:
            self.__dict__.setdefault(key, default_value)
        if value is not None:
            self.__dict__[key] = value

class NameSpace:
    def __init__(self, named_path:str, vid:str=None) -> None:
        '''命名规范, 在命名路径搜索到的文件含版本标识的文件名需要符合命名规范
        :param named_path: 将要/已经命名的路径
        :param vid: 版本id, 默认取`最新`的版本版本id且`不更新`
        '''
        self.named_path = named_path
        self._create_vid(named_path, vid)
        self.sign = generate_letter_id()

    def _create_vid(self, named_path:str, vid:str=None):
        '''创建版本id'''
        if not os.path.exists(named_path):
            os.makedirs(named_path)
        names = search_file(named_path, detectors=[
            (Detector.filename, lambda x: 'v' in x)
        ])
        names = list(set(map(lambda x:Path(x).stem.split('v')[-1], names)))
        self.update = False
        if vid:
            nvid = vid.replace('v', '')
            if nvid not in names:
                self.update = True
        else:  # 自动命名, 不稳定
            names = sorted(names) if names else [0]
            try:
                nvid = int(names[-1])
            except ValueError:
                nvid = float(names[-1])
        self.__vid = f'v{nvid}'
    
    def _search_named_file(self, vid=None):
        '''搜索命名的文件, 不稳定'''
        vid = vid or self.__vid
        logic = lambda x: vid in x
        if '.' not in vid:
            _logic = logic
            pattern = r'\w+_v\d+\.\d\w+'  #小版本
            logic = lambda x: _logic(x) and not re.match(pattern, x)  # 大版本
        return search_file(self.named_path, detectors=[
            (Detector.filename, logic)
        ])
    
    def _merge_id(self, *obj:str, has_sign:bool=False):
        '''合并id变量'''
        vid = self.get_signed_id(self.__vid) if has_sign else self.__vid
        return join_vars(*obj, vid)
    
    @property
    def vid(self):
        return self.__vid
    
    @vid.setter
    def vid(self, vid):
        self.__vid = vid

    @property
    def train_id(self):
        return self._merge_id('train')
    
    @property
    def test_id(self):
        return self._merge_id('test')
    
    def get_model_id(self, model_name:str, has_sign:bool=True):
        '''获取模型id, 无`model`前缀'''
        return self._merge_id(model_name, has_sign=has_sign)
    
    def get_signed_id(self, vid:str):
        '''获取标记id'''
        return join_vars(vid, self.sign, sep='-')
    
    @staticmethod
    def find_signed_file(signs:Sequence[str], *paths):
        '''只支持目录路径'''
        for path, sign in itertools.product(paths, signs):
            matches = glob.glob(str(
                Path(path).joinpath(f'*-{sign}.*')
            ))
            for match in matches:
                yield match

class ModelFactory(metaclass=ABCMeta):
    @abstractmethod
    def model_class(self) -> type:
        pass

    def objective(self, trial) -> object:
        return self.create_model()
    
    def create_model(self, **kwargs):
        model_class = self.model_class()
        if model_class is None:
            raise ValueError('Model class is not specified.')
        self.model = model_class(**kwargs)
        return self.model

    @property
    def model_name(self):
        return self.model_class().__name__

class RecordManager(ExitContext):
    '''记录管理'''
    TempAction = namedtuple('TempAction', ['key', 'callable', 'other_key'], defaults=[None])
    def __init__(self, save_path:str, verbose:bool=True):
        self.extend_dict = {}
        self.replenish = {}
        self.__g = {}
        self.save_path = save_path
        self.verbose = verbose
        self.fixed_report_kwargs()
    
    def fixed_report_kwargs(self):
        '''固定制作报告的参数'''
        self.report_kwargs = dict(types=['1'], indicators=['recall', 'precision'], 
                                  save_path=self.save_path, verbose=self.verbose)
    
    def _get_temp_field_name(self, field_name:str):
        '''临时字段名'''
        return f'_temp_{field_name}'

    def _temp_storage(self, field_name:str, value:Any):
        '''临时存储'''
        temp_field_name = self._get_temp_field_name(field_name)
        self.__g.setdefault(temp_field_name, [])
        self.__g[temp_field_name].append(value)
        return temp_field_name
    
    def extend_sum(self, field_name:str, num:Any):
        '''求和字段'''
        self.extend_dict[field_name] = num
        self.replenish[field_name] = self.replenish.get(field_name, 0) + num
    
    def extend_mean(self, field_name:str, num:Any):
        '''平均字段'''
        self.extend_dict[field_name] = to_percent(num)
        _callable:Callable[[List[float]], str] = lambda x: to_percent(sum(x)/len(x))

        temp_field_name = self._temp_storage(field_name, num)
        self.replenish[field_name] = self.TempAction(key=temp_field_name, callable=_callable)
    
    def extend_fixed(self, field_name:str, value:Any, storage:bool=False):
        '''固定字段'''
        self.extend_dict[field_name] = value
        self.replenish.setdefault(field_name, value)
        if storage: self._temp_storage(field_name, value)

    def extend_replenish(self, _callable:Callable[[List[float], List[float]], Union[Any]], 
                         field_name:str, other_field_name:str):
        '''自定义扩展方式'''
        temp_field_name = self._get_temp_field_name(field_name)
        other_temp_field_name = self._get_temp_field_name(other_field_name)
        self.replenish[field_name] = self.TempAction(key=temp_field_name, callable=_callable, 
                                                     other_key=other_temp_field_name)
    
    def extend_general(self, field_name:str, value:Any, storage:bool=False):
        '''普通字段'''
        self.extend_dict[field_name] = value
        self.replenish[field_name] = '-'
        if storage: self._temp_storage(field_name, value)
    
    def __setitem__(self, field_name:str, value:Any):
        self.extend_general(field_name, value)

    def _make_report(self, score:Sequence[float]=None, storage:bool=True):
        '''制作报告中转站'''
        if score is None:
            if 'score' in self.extend_dict:
                score = self.extend_dict.pop('score')
            else:
                raise ValueError('score is not found')
        if storage: self._temp_storage('score', score)
        score = list(map(lambda x:to_percent(x), score))
        auto_report(scores=score, extend_dict=self.extend_dict, **self.report_kwargs)
    
    def no_exc_clear(self):
        '''结束后清理'''
        self.replenish = {
            k: (
                v.callable(self.__g[v.key], self.__g[v.other_key])
                if v.other_key else v.callable(self.__g[v.key])
            ) if isinstance(v, self.TempAction) else v
            for k, v in self.replenish.items()
        }
        self.extend_dict = self.replenish
        self._make_report()

class AbstractTaskManager(Singleton, ExitContext, metaclass=ABCMeta):
    '''任务管理'''
    Task = namedtuple('Task', ['task', 'state'])
    def __init__(self) -> None:
        self.fields_sep = ' '
        self.values_sep = ','

    def read_tasks(self, path:str) -> list[Task]:
        '''读取任务'''
        create_path(path)
        with open(path, 'r') as f:
            return list(map(lambda x:self.Task(*x.strip().split(self.fields_sep)), f.readlines()))
        
    def save_tasks(self, tasks:list[Task], path:str) -> None:
        '''记录任务'''
        with open(path, 'w') as f:
            for _task in tasks:
                f.write(join_vars(*_task, sep=self.fields_sep) + '\n')

    def task(self, task_name:str, callback:Callable, **kwargs) -> None:
        '''任务管理流'''
        self.task_start(task_name)
        callback(self, **kwargs)
        self.task_achieve()
    
    @abstractmethod
    def task_start(self, task_name:str):
        pass

    @abstractmethod
    def task_achieve(self):
        pass

class DataManager(AbstractTaskManager):
    Task = namedtuple('Task', ['vid', 'task', 'state', 'marked_path'], defaults=[''])
    def __init__(self, log_path:str, named_path:str, vid:str=None,  
                 mode:Literal['strict', 'slack']='strict') -> None:
        '''
        :param log_path: 日志路径
        :param named_path: 将要/已经命名的路径
        :param vid: 版本id
        :param mode: 严格模式和宽松模式, 严格模式下在发生异常时会删除所有标记的路径, 
                     宽松模式下仅提示标记的路径是否存在
        '''
        super().__init__()
        self.log_path = log_path
        self.namespace = NameSpace(named_path, vid)
        self.__vid = self.namespace.vid
        self.update = self.namespace.update
        self.mode = mode
        self.marked_path = set()
        self._tasks:list[DataManager.Task] = self.read_tasks(self.log_path)

    @property
    def vid(self):
        return self.__vid
    
    @vid.setter
    def vid(self, vid):
        self.namespace.vid = vid
        self.__vid = vid
    
    def path_mark(self, root:str, vid:str=None, ext:str='', 
                  no_mark=False, has_sign:bool=False) -> Path:
        '''路径标记
        :param root: 所在文件夹
        :param vid: 版本id
        :param ext: 后缀名
        '''
        if vid is None: vid = self.__vid
        if has_sign: vid = self.namespace.get_signed_id(vid)
        path = Path(root).joinpath(vid + ext)
        if not no_mark:
            self.marked_path.add(path)
        return path
    
    def add_marked_path(self, *paths) -> None:
        '''添加标记路径'''
        self.marked_path = self.marked_path.union(set(paths))
    
    def _find_task(self, task_name:str) -> Optional[Task]:
        '''寻找任务'''
        return next(filter(lambda x:x.vid==self.__vid and x.task==task_name, self._tasks), None)

    def _update_task(self, _task:str) -> None:
        '''更新任务'''
        self._tasks[self._tasks.index(self._find_task(_task.task))] = _task
    
    def task_start(self, task_name:str) -> None:
        '''开始任务'''
        self._task = self._find_task(task_name)
        if self._task is None:
            self._task = self.Task(self.__vid, task_name, 0)
            self._tasks.append(self._task)
    
    def task_achieve(self, state:int=1) -> None:
        '''完成任务'''
        marked_path = join_vars(*self.marked_path, sep=self.values_sep)
        self._update_task(self._task._replace(state=state, marked_path=marked_path))
        self.marked_path.clear()
        self.save_tasks(self._tasks, self.log_path)
 
    def remove_version(self, vid:str, retain_task:str=None) -> None:
        '''删除特定的版本'''
        for _task in self._tasks:
            if _task.vid != vid or _task.task == retain_task:
                continue
            for path in _task.marked_path.split(self.values_sep):
                remove_path(path)
            self._update_task(_task._replace(state=0, marked_path=''))
        self.save_tasks(self._tasks, self.log_path)
    
    def rename_version(self, old_vid:str, new_vid:str):
        '''版本重命名'''
        tasks = []
        for _task in self._tasks:
            if _task.vid == old_vid:
                marked_path = set()
                for path in _task.marked_path.split(self.values_sep):
                    new_path = path.replace(old_vid, new_vid)
                    if os.path.exists(path):
                        os.rename(path, new_path)
                    marked_path.add(new_path)
                marked_path = self.values_sep.join(marked_path)
                tasks.append(_task._replace(vid=new_vid, marked_path=marked_path))
            else:
                tasks.append(_task)
        self.save_tasks(tasks, self.log_path)
        # 指向重命名后的任务
        self._tasks = tasks
            
    def exc_clear(self) -> None:
        '''异常清理'''
        paths = deepcopy(self.marked_path)
        for path in paths:
            if self.mode == 'strict':
                remove_path(path)
                self.marked_path.remove(path)
            else:
                if os.path.exists(path):
                    print(path, 'succeed')
                else:
                    self.marked_path.remove(path)
                    print(path, 'fail')
        if hasattr(self, '_task'):
            self.task_achieve(state=0)

class PresetNode:
    def __init__(self, **items) -> None:
        self.node = {}
        self.add_items(**items)
        self.children:list[PresetNode] = []
    
    @classmethod
    def named_node(cls, name, **items):
        '''命名节点'''
        obj = cls(**items)
        obj.name = name
        return obj
    
    @classmethod
    def leaf_factory(cls, keys:Sequence[Any]) -> Callable[[Sequence[Any]], PresetNode]:
        '''叶子工厂'''
        return lambda values:cls(**dict(zip(keys, values)))
    
    def add_items2(self, keys, values):
        '''添加元素'''
        for key, value in zip(keys, values):
            self.node[key] = value

    def add_items(self, **items):
        '''添加元素'''
        for key, value in items.items():
            self.node[key] = value
    
    def add_child(self, *children):
        '''添加子节点'''
        self.children.extend(children)
        return self
    
    def mount(self):
        '''挂载叶子节点'''
        if hasattr(self, 'name'):
            self.node[self.name] = self.children
            self.children = []
        return self

    def to_dict(self):
        '''转字典形式'''
        root = deepcopy(self.node)
        for child in self.children:
            if hasattr(child, 'name'):
                root[child.name] = [child.to_dict() for child in child.node.get(child.name, [])]
        return root
    
    @staticmethod
    def sync_item(obj:Dict, other:Dict, *keys:Any, 
                  map_keys:Tuple[Sequence[Any], Sequence[Any]]=None) -> Optional[Any]:
        '''同步元素'''
        if len(keys) > 0:
            value = PresetNode.get_value(obj, *keys)
            PresetNode.set_value(other, value, *keys)
            return value
        if map_keys:
            value = PresetNode.get_value(obj, *map_keys[0])
            PresetNode.set_value(other, value, *map_keys[1])
            return value
    
    @staticmethod
    def set_value(obj:Dict, value:Any, *keys):
        cur = obj
        for key in keys[:-1]:
            cur[key] = cur.get(key, {})
            cur = cur[key]
        cur[keys[-1]] = value
    
    @staticmethod
    def get_value(obj:Dict, *keys) -> Any:
        '''查值'''
        cur = obj
        for key in keys:
            if cur is None:
                break
            cur = cur.get(key, None)
        return cur
    
class SearchSpace:
    '''搜索空间'''
    def __init__(self, preset:Dict[str, Any]) -> None:
        self.preset = preset
        self.space_division(self.preset)
        
    def space_division(self, *objs:Dict[str, Any], keys:List[str]=None):
        '''划分空间'''
        keys = keys or ['p', 'desc']
        for obj, key in itertools.product(objs, keys):
            obj.setdefault(key, {})
    
    def get_p(self, key:str):
        '''获取指针'''
        self.preset['p'][key] = self.preset['p'].get(key, 0)
        return self.preset['p'][key]
    
    def set_p(self, key:str, p:int):
        '''设置指针'''
        self.preset['p'][key] = p
    
    def record_desc(self, key:str, item:Any):
        '''记录描述信息'''
        desc_lis:List[Any] = self.preset['desc'].get(key, [])
        desc_lis.append(str(item))
        self.preset['desc'][key] = desc_lis

    def make_preset(self, key:str, items:Sequence[Any]) -> None:
        '''下一个'''
        for item in items:
            p = self.get_p(key)
            self.set_p(key, p + 1)  # 存下一个的位置
            self.record_desc(key, item)
    
class SearchSpaceManager(SearchSpace, Singleton, Functional, ExitContext):
    '''预处理'''
    def __init__(self, preset_path:str, log_path:str, 
                 skip_selections:List[int]=None, 
                 skip_mode:Optional[Literal['follow']]=None) -> None:
        self.preset_path = preset_path
        self.log_path = log_path
        
        self.selections = {}
        self.selection_freq = 0
        self.skip_selections = skip_selections or []
        self.skip_mode = skip_mode

        self.old_selections: Dict[str, Any] = self.get_old_selections(log_path)
        self.preset: Dict[str, Any] = loadJSON(preset_path)
        self.space_division(self.preset, self.selections)
    
    def get_old_selections(self, path):
        producer = lambda x: re.sub(r'-\w+', '*', str(x))  # 搜索范围
        signed_path = get_latest_modified_file(producer(path))
        if signed_path: return loadJSON(signed_path)
        return {}
            
    def clear(self):
        '''保存参数信息'''
        writeJSON(self.preset, self.preset_path)
        writeJSON(self.selections, self.log_path)
    
    def get_query_chain(self, keys, p:bool=True):
        '''查询链'''
        assert len(keys) > 0
        names = ['p', 'desc']  # 块名
        if not p: return [None] * len(names)
        if not isinstance(keys, list):
            keys:List = list(keys)
        chains = [deepcopy(keys) for _ in names]
        for chain, name in zip(chains, names):
            chain.insert(0, name)
        return chains
    
    @print_freq('<<<<<<<<< {} >>>>>>>>>')
    def load_preset(self, *keys:str, values:Optional[Sequence[Any]]=None) -> Optional[Dict[str, str]]:
        '''预设和加载配置信息'''
        pkeys, desckeys = self.get_query_chain(keys, p=values)
        selection = None
        self.selection_freq += 1
        if self.selection_freq in self.skip_selections:
            if self.skip_mode == 'follow':
                if values:
                    selection = PresetNode.sync_item(self.old_selections, self.selections, *pkeys)
                else:
                    selection = PresetNode.sync_item(self.old_selections, self.selections, *keys)
        if selection is None:
            if values:
                lis = PresetNode.get_value(self.preset, *desckeys)
                if lis:
                    desc, index = custom_selection(lis, return_index=True)
                    PresetNode.set_value(self.selections, index, *pkeys)
                    PresetNode.set_value(self.selections, desc, *desckeys)
                    selection = values[index]
            else:
                lis = PresetNode.get_value(self.preset, *keys)
                if lis:
                    selection = custom_selection(lis)
                    PresetNode.set_value(self.selections, selection, *keys)
        return selection
        
    def make_preset(self, key:str, values:Sequence[Any]=None,
                    subkeys:Sequence[str]=None, 
                    subvalues:Sequence[Any]=None,
                    multi_subvalues:Sequence[Sequence[Any]]=None):
        '''两种快速制作预设数据的方法'''
        if values:  # 对象
            super().make_preset(key, values)
        if subkeys:  # 字符串或数字
            root = PresetNode()
            main_node = PresetNode.named_node(key)
            if subvalues:
                main_node.add_items2(subkeys, subvalues)
            if multi_subvalues:
                leaf_factory = PresetNode.leaf_factory(subkeys)
                for subvalues in multi_subvalues:
                    leaf = leaf_factory(subvalues)
                    main_node.add_child(leaf)
                main_node.mount()
            preset = root.add_child(main_node).to_dict()
            self.preset = {**self.preset, **preset}
        return self.preset

class TestAlgorithm:
    @staticmethod
    def one_class_intersection(y_true:Sequence, y_pred:Sequence, class_:int=1, 
                               fill:bool=False) -> Tuple[float, float]:
        '''单类交集'''
        ttrue, ptrue, tnum, pnum = 0, 0, 0, 0
        temp, flag = [], False
        for n, (i, j) in enumerate(zip(y_true, y_pred)):
            if i == class_:
                if i == j:
                    flag = True
                    if fill:
                        ptrue += 1
                temp.append(n)

                if n+1 < len(y_true) and y_true[n+1] != class_:
                    tnum += 1  # 成块, 多窗口映射重合归属问题
                    if flag:
                        ttrue += 1
                        if not fill:
                            ptrue += len(temp)
                    temp, flag = [], False
                    
            if j == class_:
                pnum += 1
        return ttrue/tnum, ptrue/pnum
    
    @staticmethod
    def merge_similar_events(y_true:Sequence, y_pred:Sequence) -> Tuple[Sequence, Sequence]:
        '''合并同类事件'''
        _y_true, _y_pred, temp, cat = [], [], [], y_true[0]
        for i in range(len(y_true)):
            if cat == y_true[i]:
                temp.append(y_pred[i])
            elif cat != y_true[i] or i == len(y_true)-1:
                _y_true.append(cat)
                if cat in temp:
                    _y_pred.append(cat)
                else:
                    val = Counter(temp).most_common(1)[0][0]
                    _y_pred.append(val)
                cat, temp = y_true[i], [y_pred[i]]
            else:
                _y_true.append(y_true[i])
                _y_pred.append(y_pred[i])
        return _y_true, _y_pred
    
    @staticmethod
    def virtual_tagscmp2(src:Sequence[Tuple[int, int]], dst:Sequence[Tuple[int, int]], 
                        threshold:int=0) -> Tuple[float, float]:
        '''缩小预测区域'''
    
    @staticmethod
    def virtual_tagscmp(src:Sequence[Tuple[int, int]], dst:Sequence[Tuple[int, int]], 
                        threshold:int=0) -> Tuple[float, float]:
        '''虚拟tags比较, 适用于单文件单标签'''
        ttrue, ptrue = 0, 0
        for i in src:
            scores = []
            for j in dst:
                sorce = jaccard_similarity(i, j)
                scores.append(sorce)
            temp = [s for s in scores if s > threshold]
            if sum(temp) > threshold:
                ttrue += 1
                ptrue += len(temp)
        tiny = 1e-6
        tnum = len(src) if src else tiny
        pnum = len(dst) if dst else tiny
        return ttrue/tnum, ptrue/pnum
    
    @staticmethod
    def weighted_average(scores:list[Union[Sequence[float], float]], 
                         weights:list[Union[int, float]]) -> Union[list[np.ndarray], np.ndarray]:
        '''加权平均'''
        if isinstance(scores[0], Sequence):
            return np.sum([weights[i]/sum(weights)*np.array(scores[i]) for i in range(len(scores))], axis=0)
        return np.sum(np.array(weights)/sum(weights)*np.array(scores))
